#include "triu_custom_tiling.h"
#include "kernel_operator.h"


constexpr int32_t BUFFER_NUM = 2;


class KernelTriuCustom
{
public:
    __aicore__ inline KernelTriuCustom() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, TriuCustomTilingData tiling)
    {
        this->matrixNum = tiling.matrixNum;
        this->rowLength = tiling.rowLength;
        this->columnLength = tiling.columnLength;
        this->diagVal = tiling.diagVal;

        uint32_t singleBuffer = tiling.fullTileLength;
        if (singleBuffer < tiling.lastTileLength)
        {
            singleBuffer = tiling.lastTileLength;
        }

        xGm.SetGlobalBuffer((__gm__ half *)x, tiling.totalLengthAligned);
        yGm.SetGlobalBuffer((__gm__ half *)y, tiling.totalLengthAligned);
        pipe.InitBuffer(queBind, BUFFER_NUM, singleBuffer * sizeof(half));
    }
    __aicore__ inline void Process()
    {
        int32_t cnt = 0;
        for (int32_t i = 0; i < this->matrixNum; i++)
        {
            for (int32_t j = 0; j < this->rowLength; j++)
            {
                int32_t k = 0;
                while (k < this->columnLength && k - j < this->diagVal)
                {
                    yGm.SetValue(cnt, (half)0);
                    k++;
                    cnt++;
                }
                while (k < this->columnLength)
                {
                    half curr = xGm.GetValue(cnt);
                    yGm.SetValue(cnt, curr);
                    k++;
                    cnt++;
                }
            }
        }
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQueBind<AscendC::QuePosition::VECIN, AscendC::QuePosition::VECOUT, BUFFER_NUM> queBind;
    AscendC::GlobalTensor<half> xGm;
    AscendC::GlobalTensor<half> yGm;

    int32_t matrixNum;
    int32_t rowLength;
    int32_t columnLength;
    int32_t diagVal;

};

extern "C" __global__ __aicore__ void triu_custom(GM_ADDR x, GM_ADDR y, TriuCustomTilingData tiling)
{
    KernelTriuCustom op;
    op.Init(x, y, tiling);
    op.Process();
}
